import torch
import json
import os
from torch.utils.data import DataLoader
from tqdm import tqdm
from sklearn.metrics import precision_recall_fscore_support
from config import Config
from model import BertNLU
from dataset import NLUDataset
from transformers import BertTokenizer
from nlp_model.logger import get_logger

# 初始化日志
logger = get_logger("predict")


def load_resources():
    """加载模型、映射表和分词器（保持原有逻辑）"""
    with open(f"{Config.model_dir}/intent_map.json", 'r', encoding='utf-8') as f:
        intent_map = json.load(f)
    with open(f"{Config.model_dir}/slot_map.json", 'r', encoding='utf-8') as f:
        slot_map = json.load(f)

    tokenizer = BertTokenizer.from_pretrained(Config.bert_model)

    model = BertNLU(
        num_intents=len(intent_map['intent2id']),
        num_slots=len(slot_map['slot2id'])
    ).to(Config.device)
    model_path = os.path.join(Config.model_dir, "bert_model.bin")
    model.load_state_dict(torch.load(model_path, map_location=Config.device))
    model.eval()

    return model, intent_map, slot_map, tokenizer


def load_test_data(test_path):
    """整体读取test.json（适配“单文件大JSON对象”格式）"""
    try:
        # 一次性读取整个文件，解析为完整JSON对象
        with open(test_path, 'r', encoding='utf-8') as f:
            test_data = json.load(f)
        # 提取所有NLUxxxx样本，转为列表（适配NLUDataset的输入格式）
        # 假设NLUDataset需要的样本格式是：[{"text": "...", "intent": "...", "slots": {...}}, ...]
        samples = [sample for sample in test_data.values()]
        logger.info(f"测试集加载成功 | 共 {len(samples)} 条数据")
        return samples
    except json.JSONDecodeError as e:
        logger.error(f"test.json格式错误: {str(e)}")
        raise
    except Exception as e:
        logger.error(f"测试集读取失败: {str(e)}")
        raise


def update_slot_map(slot_map, test_samples):
    """根据测试集的槽位名称，补充映射表中缺失的槽位标签（生成BIO标签）"""
    slot2id = slot_map['slot2id'].copy()
    id2slot = slot_map['id2slot'].copy()
    # 从测试集slots中提取所有槽位名称，生成对应的B-xxx和I-xxx标签
    test_slot_names = set()
    for sample in test_samples:
        test_slot_names.update(sample['slots'].keys())

    # 生成需要的BIO标签（如appliance→B-appliance、I-appliance）
    required_labels = {"O"}  # 必须包含O标签
    for slot_name in test_slot_names:
        required_labels.add(f"B-{slot_name}")
        required_labels.add(f"I-{slot_name}")

    # 补充缺失的标签到映射表
    existing_labels = set(slot2id.keys())
    missing_labels = required_labels - existing_labels
    if missing_labels:
        logger.warning(f"检测到测试集缺失的槽位标签：{missing_labels}，自动补充到映射表")
        max_id = max(slot2id.values()) if slot2id else 0
        for label in missing_labels:
            max_id += 1
            slot2id[label] = max_id
            id2slot[str(max_id)] = label  # 适配id2slot的字符串键格式
        slot_map['slot2id'] = slot2id
        slot_map['id2slot'] = id2slot
    return slot_map


def evaluate_test_set():
    # 1. 加载基础资源
    model, intent_map, slot_map, tokenizer = load_resources()

    # 2. 正确读取测试集（整体解析）
    test_samples = load_test_data(Config.dev_path)

    # 3. 根据测试集槽位名称，补充BIO标签到映射表
    slot_map = update_slot_map(slot_map, test_samples)

    # 4. 获取O标签ID（用于过滤）
    o_label_id = slot_map['slot2id'].get('O', -1)
    if o_label_id == -1:
        logger.warning("槽位映射表中未找到'O'标签，将不进行过滤")

    # 5. 手动构造NLUDataset所需的格式（确保包含生成标签的必要信息）
    # 临时保存处理后的测试集到内存（不修改原文件），供NLUDataset读取
    # （若NLUDataset支持直接接收样本列表，可跳过此步，需根据dataset.py逻辑调整）
    # 这里假设NLUDataset需要从文件读取，因此临时生成一个符合格式的JSON文件
    temp_test_path = os.path.join(Config.model_dir, "temp_test.json")
    with open(temp_test_path, 'w', encoding='utf-8') as f:
        json.dump(test_samples, f, ensure_ascii=False, indent=2)

    # 6. 加载测试集（使用临时文件，确保格式适配）
    test_set = NLUDataset(temp_test_path, tokenizer, is_test=False)
    test_set.intent2id, test_set.id2intent = intent_map['intent2id'], intent_map['id2intent']
    test_set.slot2id, test_set.id2slot = slot_map['slot2id'], slot_map['id2slot']
    test_loader = DataLoader(
        test_set,
        batch_size=Config.batch_size,
        shuffle=False,
        num_workers=2
    )
    logger.info(f"测试集Dataset加载完成 | 共 {len(test_set)} 条数据")

    # 7. 预测与评估（保持原有逻辑）
    all_intent_preds = []
    all_intent_true = []
    all_slot_preds = []
    all_slot_true = []

    with torch.no_grad():
        for batch in tqdm(test_loader, desc="测试集预测中"):
            input_ids = batch['input_ids'].to(Config.device)
            attention_mask = batch['attention_mask'].to(Config.device)
            intent_ids = batch['intent_id'].to(Config.device)
            slot_ids = batch['slot_ids'].to(Config.device)

            intent_logits, slot_logits = model(input_ids, attention_mask)
            intent_preds = intent_logits.argmax(1)
            slot_preds = slot_logits.argmax(2)

            all_intent_preds.extend(intent_preds.cpu().numpy())
            all_intent_true.extend(intent_ids.cpu().numpy())

            mask = attention_mask.bool()
            valid_mask = (slot_ids != o_label_id) & mask if o_label_id != -1 else mask
            all_slot_preds.extend(slot_preds[valid_mask].cpu().numpy())
            all_slot_true.extend(slot_ids[valid_mask].cpu().numpy())

    # 8. 计算指标（保持原有逻辑）
    intent_acc = sum(p == t for p, t in zip(all_intent_preds, all_intent_true)) / len(all_intent_true)
    intent_precision, intent_recall, intent_f1, _ = precision_recall_fscore_support(
        all_intent_true, all_intent_preds, average='macro', zero_division=0
    )

    slot_metrics = (0.0, 0.0, 0.0, 0.0)
    if len(all_slot_true) > 0:
        slot_acc = sum(p == t for p, t in zip(all_slot_preds, all_slot_true)) / len(all_slot_true)
        slot_precision, slot_recall, slot_f1, _ = precision_recall_fscore_support(
            all_slot_true, all_slot_preds, average='macro', zero_division=0
        )
        slot_metrics = (slot_acc, slot_precision, slot_recall, slot_f1)
    else:
        slot_acc = slot_precision = slot_recall = slot_f1 = 0.0

    # 9. 日志输出
    logger.info("\n" + "=" * 100)
    logger.info(f"测试集评估结果（共 {len(test_set)} 条数据）")
    logger.info("-" * 100)
    logger.info("【意图识别指标】")
    logger.info(f"  准确率（Accuracy）: {intent_acc:.4f}")
    logger.info(f"  精确率（Precision）: {intent_precision:.4f}")
    logger.info(f"  召回率（Recall）: {intent_recall:.4f}")
    logger.info(f"  F1值（F1-Score）: {intent_f1:.4f}")
    logger.info("-" * 100)
    logger.info("【槽位填充指标】（已排除'O'标签，仅统计有效槽位）")
    logger.info(f"  准确率（Accuracy）: {slot_acc:.4f}")
    logger.info(f"  精确率（Precision）: {slot_precision:.4f}")
    logger.info(f"  召回率（Recall）: {slot_recall:.4f}")
    logger.info(f"  F1值（F1-Score）: {slot_f1:.4f}")
    logger.info("=" * 100 + "\n")

    # 10. 删除临时文件
    os.remove(temp_test_path)


if __name__ == "__main__":
    evaluate_test_set()